from google.colab import files

uploaded = files.upload()

for fn in uploaded.keys():
  print('User uploaded file "{name}" with length {length} bytes'.format(
      name=fn, length=len(uploaded[fn])))
  with open(fn, 'wb') as f:
    f.write(uploaded[fn])
  print('File uploaded successfully!')

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
tokenizer.pad_token = tokenizer.eos_token

# Tokenize dataset
def tokenize_function(examples):
    tokens = tokenizer(examples["text"], padding="max_length", truncation=True)
    tokens["labels"] = tokens["input_ids"].copy()  # Labels are the same as input_ids for language modeling
    return tokens

tokenized_datasets = dataset.map(tokenize_function, batched=True)
tokenized_datasets

# Load model for causal language modeling
model = AutoModelForCausalLM.from_pretrained(checkpoint)

# Training arguments
training_args = TrainingArguments(
    output_dir="./results",
    learning_rate=2e-4,
    per_device_train_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets,
    tokenizer=tokenizer,
)

# Train the model
trainer.train()

# Generate encoded text
with torch.no_grad():
    outputs = model.generate(input_ids, min_length=100, max_length=500, num_return_sequences=1, temperature=0.7, top_k=50)

# Decode and print the generated text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)